from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pathlib import Path
import yaml
from loguru import logger
import sys
import os
from fastapi.middleware.cors import CORSMiddleware
import json
from datetime import datetime
from typing import List, Dict, Any
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 添加项目根目录到 Python 路径
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(ROOT_DIR)

from script.osv.collector import OSVCollector

app = FastAPI(title="CSA Anti-Poison Scanning Engine")

# 配置 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173"],  # 允许前端域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头部
)

# 初始化收集器
collector = OSVCollector()

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory=os.path.join(ROOT_DIR, "frontend/static")), name="static")

@app.get("/")
async def home(page: int = 1, page_size: int = 10):
    """首页"""
    try:
        # 获取统计数据
        stats = collector.get_stats()
        # 获取最近的漏洞（带分页）
        recent_vulns = collector.get_recent_vulnerabilities(page_size, page)
        # 获取总漏洞数用于计算总页数
        total_vulns = collector.get_total_vulnerabilities()
        
        return JSONResponse(content={
            "stats": stats,
            "recent_vulnerabilities": recent_vulns,
            "pagination": {
                "current_page": page,
                "page_size": page_size,
                "total_items": total_vulns,
                "total_pages": (total_vulns + page_size - 1) // page_size
            }
        })
    except Exception as e:
        logger.error(f"Error getting home data: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/{path:path}")
async def catch_all(path: str):
    """处理所有前端路由"""
    return FileResponse(os.path.join(ROOT_DIR, 'frontend/index.html'))

@app.get("/api/vulnerabilities")
async def get_all_vulnerabilities():
    """获取所有漏洞"""
    try:
        vulnerabilities = collector.get_all_vulnerabilities()
        return JSONResponse(content={"vulnerabilities": vulnerabilities})
    except Exception as e:
        logger.error(f"Error getting all vulnerabilities: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/vulnerabilities/recent")
async def get_recent_vulnerabilities(limit: int = 10):
    """获取最近的漏洞"""
    try:
        vulnerabilities = collector.get_recent_vulnerabilities(limit)
        return JSONResponse(content={"vulnerabilities": vulnerabilities})
    except Exception as e:
        logger.error(f"Error getting recent vulnerabilities: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/vulnerabilities/search")
async def search_vulnerabilities(query: str):
    """搜索漏洞"""
    try:
        results = collector.search_vulnerabilities(query)
        return JSONResponse(content={"vulnerabilities": results})
    except Exception as e:
        logger.error(f"Error searching vulnerabilities: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/samples")
async def get_samples():
    """获取样本数据"""
    try:
        samples = collector.get_samples()
        return JSONResponse(content={"samples": samples})
    except Exception as e:
        logger.error(f"Error getting samples: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/stats")
async def get_stats():
    """获取漏洞统计信息"""
    try:
        stats = collector.get_stats()
        return JSONResponse(content=stats)
    except Exception as e:
        logger.error(f"Error getting stats: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001) 